package com.northpool.operator.statistics.dataset.helper;

public class LogisticRegression {
    private double[] weights;
    private double learningRate;
    private int iterations;

    public LogisticRegression(double learningRate, int iterations) {
        this.learningRate = learningRate;
        this.iterations = iterations;
    }

    // Sigmoid函数
    private double sigmoid(double z) {
        return 1 / (1 + Math.exp(-z));
    }

    // 训练模型
    public void fit(double[][] X, int[] y) {
        int m = X.length;
        int n = X[0].length;
        weights = new double[n];

        for (int i = 0; i < iterations; i++) {
            double[] gradients = new double[n];
            for (int j = 0; j < m; j++) {
                double predicted = predictProb(X[j]);
                for (int k = 0; k < n; k++) {
                    gradients[k] += (predicted - y[j]) * X[j][k];
                }
            }
            for (int k = 0; k < n; k++) {
                weights[k] -= learningRate * gradients[k] / m;
            }
        }
    }

    // 预测概率
    public double predictProb(double[] x) {
        double z = 0.0;
        for (int i = 0; i < weights.length; i++) {
            z += weights[i] * x[i];
        }
        return sigmoid(z);
    }

    // 预测类别
    public int predict(double[] x) {
        return predictProb(x) >= 0.5 ? 1 : 0;
    }

    public static void main(String[] args) {
        // 示例数据：数学和物理成绩
        double[][] X = {
                {55, 60}, {70, 65}, {80, 75}, {90, 85},
                {50, 55}, {40, 45}, {60, 50}, {75, 80},
                {85, 70}, {95, 90}, {45, 40}, {65, 55},
                {85, 80}, {70, 75}, {80, 85}
        };
        int[] y = {0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 1}; // 0: 不及格, 1: 及格

        LogisticRegression model = new LogisticRegression(0.01, 1000);
        model.fit(X, y);

        // 预测新的学生成绩
        double[] newStudent = {70, 75}; // 数学和物理成绩
        int result = model.predict(newStudent);
        System.out.println("预测结果: " + (result == 1 ? "及格" : "不及格"));
    }
}
